'''
Created on May 21, 2012

@author: tel
'''
import os, sys
import numpy as np
import matplotlib
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.ticker import LogFormatter 
from matplotlib.ticker import MaxNLocator
from pylab import setp
from plothelper import Unpickle, Repickle, Polyfit

class LogFormatterHB(LogFormatter): 
    def __call__(self, v, pos=None): 
        vv = self._base ** v 
        return LogFormatter.__call__(self, vv, pos) 

def OffTicks(ax, list):
    for entry in list:
        if entry=='x':
            setp(ax.get_xticklabels(), visible=False)
        elif entry=='y':
            setp(ax.get_yticklabels(), visible=False)

def Hexbin(*args, **kwargs):
    if 'outfile' in kwargs:
        outfile = kwargs['outfile']
    else:
        outfile = 'Hexbin'
    if 'outdir' in kwargs:
        outdir = kwargs['outdir']
    else:
        outdir = '/home/tel/benchmark/links'
    if 'repickle' in kwargs:
        repickle = kwargs['repickle']
    else:
        repickle = True
    x,y = args[0],args[1]

    if outfile=='dsol_ele_hexbin':
        xlabel = '$\Delta G_{dsol}$'
        ylabel = '$\Delta G_{ele}$'
    elif outfile=='ratio_ddpk_hexbin':
        xlabel = r'$\frac{\Delta {G}_{ele}}{\Delta {G}_{dsol}}$'
        ylabel = '$\Delta\Delta pK_{a}$'
    elif outfile=='size_deepcount_hexbin':
        xlabel = r'Protein size'
        ylabel = r'Highly buried\nresidue count'

    if repickle:
        Repickle(outdir, outfile, *args)
    
    fig = plt.figure(1)
    axes = plt.axes()
    
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)

    for tick in axes.xaxis.get_major_ticks():
        tick.label.set_fontsize('small')
        tick.label.set_rotation(45)
    for tick in axes.yaxis.get_major_ticks():
        tick.label.set_fontsize('small')
    
    if outfile=='dsol_ele_hexbin':
        hexb = axes.hexbin(x,y, cmap=cm.jet, gridsize=125, bins='log', mincnt=1)
        axes.axvline(x=6, linestyle='--')
        axes.axhline(y=0, linestyle='--')
        axes.set_ylim(-40,40)
        axes.set_xlim(0,12)
        cb = fig.colorbar(hexb, ticks=MaxNLocator(integer=True), format=LogFormatterHB())
    elif outfile=='ratio_ddpk_hexbin':
        hexb = axes.hexbin(x,y, cmap=cm.jet, gridsize=50, mincnt=1)
        cb = fig.colorbar(hexb, ticks=MaxNLocator(integer=True))
        axes.set_ylim(-20,15)
        axes.set_xlim(-4,2.5)
    elif outfile=='size_deepcount_hexbin':
        hexb = axes.hexbin(x,y, cmap=cm.jet, gridsize=50, mincnt=1)
        cb = fig.colorbar(hexb, ticks=MaxNLocator(integer=True))
        axes.set_ylim(0,65)
        axes.set_xlim(0,1000)
    else:
        hexb = axes.hexbin(x,y, cmap=cm.jet, gridsize=50, mincnt=1)
        cb = fig.colorbar(hexb, ticks=MaxNLocator(integer=True))
    pfit = Polyfit(x,y)
    fx = [-100,1000]
    fy = []
    for ifx in fx:
        fy.append(pfit['poly'][0]*ifx+pfit['poly'][1])
    axes.plot(fx,fy)
    print pfit['det']
    axes.text(.05,.95, 'm = %3.3f, b = %3.3f, R^2 = %3.3f' % (pfit['poly'][0],pfit['poly'][1],pfit['det']),horizontalalignment='left',verticalalignment='top',transform=axes.transAxes)
    cb.set_label('count')
    
    fig.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight', transparent=True)
    fig.clf()

def MultiHexbin(*args, **kwargs):
    if 'outfile' in kwargs:
        outfile = kwargs['outfile']
    else:
        outfile = 'MultiHexbin'
    if 'outdir' in kwargs:
        outdir = kwargs['outdir']
    else:
        outdir = '/home/tel/benchmark/links'
    if 'repickle' in kwargs:
        repickle = kwargs['repickle']
    else:
        repickle = True

    if repickle:
        Repickle(outdir, outfile, *args)
    
    fig = plt.figure(figsize=(10*len(args)/2,10))
    gs = gridspec.GridSpec(1,len(args)/2)
    gs.update(wspace=0.04, hspace=0.07)
    axes = []
    
    axparams = {'dsol':{'lim':(0,12)},
                'ele':{'lim':(-40,40)},
                'ratio':{'lim':(-6,5)},
                'ddpk':{'lim':(-20,15)}}
    setparams = {('dsol_ele','alldsol', 'every'):{'hexbkwargs':{'bins':'log','gridsize':125,'extent':(0,12,-40,40)},
                                                  'cbkwargs':{'format':LogFormatterHB()}},
                 ('ratio_ddpk','highdsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-6,5,-20,15)},
                                                     'cbkwargs':{}}}
    rawtags = outfile.split('.')[-1].split('-')
    settags = []
    subnames = []
    for dsol in rawtags[1].split('_'):
        settag = (rawtags[0],dsol,'every')
        if settag in setparams:
            settags.append(settag)
        else:
            for s in rawtags[2].split('_'):
                settags.append((rawtags[0],dsol,s))
    xtag = rawtags[0].split('_')[0]
    ytag = rawtags[0].split('_')[1]
    
    for i,(x,y) in enumerate(zip(args[::2],args[1::2])):
        ax = plt.subplot(gs[0, i])
        axes.append(ax)
        if outfile=='e4_dsol_e20_dsol_multihexbin':
            extent = ((-5,5,-5,5),(0,10,-5,5),(5,15,-5,5))[i]
        elif outfile=='e4_ele_e20_ele_multihexbin':
            extent = ((-15,10,-20,5),(-20,5,-20,5),(-25,0,-20,5))[i]
        elif outfile=='e4_ratio_e20_ratio_multihexbin':
            extent = ((-30,15,-30,15),(-15,5,-15,5),(-10,5,-10,5))[i]
        hexb = ax.hexbin(x,y, cmap=cm.jet, mincnt=.001, **setparams[settag]['hexbkwargs'])
        cb = fig.colorbar(hexb, ticks=MaxNLocator(integer=True), **setparams[settag]['cbkwargs'])
        ax.set_xlim(axparams[xtag]['lim'])
        ax.set_ylim(axparams[ytag]['lim'])
        pfit = Polyfit(x,y)
        fx = [-100,100]
        fy = []
        for ifx in fx:
            fy.append(pfit['poly'][0]*ifx+pfit['poly'][1])
        ax.plot(fx,fy)
        print pfit['det']
        ax.text(.05,.95, 'm = %3.3f, b = %3.3f, R^2 = %3.3f' % (pfit['poly'][0],pfit['poly'][1],pfit['det']),horizontalalignment='left',verticalalignment='top',transform=ax.transAxes)
        
    for ax in axes:
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            tick.label.set_rotation(45)
        for tick in ax.yaxis.get_major_ticks():
            tick.label.set_fontsize('small')
    
    d,delta= None,None
    if outfile=='quick_dpk_full_dpk_multihexbin':
        d = 'd'
        delta = r'\Delta '
    elif outfile=='quick_ddpk_full_ddpk_multihexbin':
        d = 'dd'
        delta = r'\Delta\Delta '
    if d!=None and delta!=None:
        xlabels = [r'Quick $%spK_{a}$' % delta]*3
        ylabels = [r'Full $%spK_{a}$' % delta]*3
        titles = ['low dsol', 'mid dsol', 'high dsol']
        for i, ax in enumerate(axes):
            ax.set_xlabel(xlabels[i])
            ax.set_ylabel(ylabels[i])
            ax.set_title(titles[i])
            ax.plot([-15,25],[-15,25])
    
    elem = None        
    if outfile=='e4_dsol_e20_dsol_multihexbin':
        elem='dsol'
        xlims = ((-5,5),(0,10),(5,15))
        ylims = ((-5,5),(-5,5),(-5,5))
    elif outfile=='e4_ele_e20_ele_multihexbin':
        elem='ele'
        xlims = ((-15,10),(-20,5),(-25,0))
        ylims = ((-20,5),(-20,5),(-20,5))
    elif outfile=='e4_ratio_e20_ratio_multihexbin':
        elem='ratio'
        xlims = ((-30,15),(-15,5),(-10,5))
        ylims = ((-30,15),(-15,5),(-10,5))
    if elem!= None:
        xlabels = [r'%s ${\epsilon}_{p}=4$' % elem]*3
        ylabels = [r'%s ${\epsilon}_{p}=20$' % elem]*3
        titles = ['low dsol', 'mid dsol', 'high dsol']
        for i, ax in enumerate(axes):
            ax.set_xlabel(xlabels[i])
            ax.set_ylabel(ylabels[i])
            ax.set_title(titles[i])
            ax.set_xlim(xlims[i])
            ax.set_ylim(ylims[i])
            ax.plot([-100,100],[-100,100])
    else:
        xlables = []
        
    plt.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight', transparent=True)
    plt.clf()
    
def MultiHexbinS(*args, **kwargs):
    if 'outfile' in kwargs:
        outfile = kwargs['outfile']
    else:
        outfile = 'MultiHexbin'
    if 'outdir' in kwargs:
        outdir = kwargs['outdir']
    else:
        outdir = '/home/tel/benchmark/links'
    if 'repickle' in kwargs:
        repickle = kwargs['repickle']
    else:
        repickle = True

    if repickle:
        Repickle(outdir, outfile, *args)
    
    axparams = {'dsol':{'lim':(0,12)},
                'ele':{'lim':(-40,40)},
                'ratio':{'lim':(-6,5)},
                'ddpk':{'lim':(-20,15)},
                'dpk':{'lim':(-15,20)},
                'dpke4':{'lim':(-15,20)},
                'dpke20':{'lim':(-15,20)},
                'dpkq':{'lim':(-20,15)},
                'dpkf':{'lim':(-20,15)},
                'ddpkq':{'lim':(-15,20)},
                'ddpkf':{'lim':(-15,20)},
                'eleq':{'lim':(-15,20)},
                'elef':{'lim':(-15,20)},
                'dsolq':{'lim':(-15,20)},
                'dsolf':{'lim':(-15,20)},
                'neujigglef':{'lim':(0,10)},
                'neupkjigglef':{'lim':(0,10)},
                'ionjigglef':{'lim':(0,10)},
                'ionpkjigglef':{'lim':(0,10)},
                'eleanom':{'lim':(-10,10)},
                'dsolanom':{'lim':(-10,10)}}
                    
    setparams = {('dsol_ele','alldsol', 'every'):{'hexbkwargs':{'bins':'log','gridsize':125,'extent':(0,12,-40,40)},
                                                  'cbkwargs':{'format':LogFormatterHB()}},
                 ('ratio_ddpk','highdsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-6,5,-20,15)},
                                                     'cbkwargs':{}},
                 ('dpk_ddpk','alldsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-15,20,-20,15)},
                                                     'cbkwargs':{}},
                 ('dpke4_dpke20','alldsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-15,20,-15,20)},
                                                     'cbkwargs':{}},
                 ('dpkq_dpkf','alldsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-20,15,-20,15)},
                                                     'cbkwargs':{}},
                 ('ddpkq_ddpkf','alldsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-15,20,-15,20)},
                                                     'cbkwargs':{}},
                 ('eleq_elef','alldsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-15,20,-15,20)},
                                                     'cbkwargs':{}},
                 ('dsolq_dsolf','alldsol', 'every'):{'hexbkwargs':{'gridsize':125,'extent':(-15,20,-15,20)},
                                                     'cbkwargs':{}}}
    figx = 12
    figy = 10
    fig = plt.figure(figsize=(figx*len(args)/2,figy))
    gs = gridspec.GridSpec(1,len(args)/2)
    gs.update(wspace=0.15, hspace=0.07)
    axes = []
    rawtags = outfile.split('.')[-1].split('-')
    settags = []
    for dsol in rawtags[1].split('_'):
        settag = (rawtags[0],dsol,'every')
        if settag in setparams.keys():
            settags.append(settag)
        else:
            for s in rawtags[2].split('_'):
                settags.append((rawtags[0],dsol,s))
                
    xtag = rawtags[0].split('_')[0]
    ytag = rawtags[0].split('_')[1]
    
    for i,(x,y) in enumerate(zip(args[::2],args[1::2])):
        ax = plt.subplot(gs[0, i])
        axes.append(ax)
        try:
            hexbkwargs = setparams[settag]['hexbkwargs']
        except KeyError:
            hexbkwargs = {'gridsize':125,'extent':axparams[xtag]['lim']+axparams[ytag]['lim']}
        hexb = ax.hexbin(x,y, cmap=cm.jet, mincnt=.001, **hexbkwargs)
        try:
            cbkwargs = setparams[settag]['cbkwargs']
        except KeyError:
            cbkwargs = {}
        cb = fig.colorbar(hexb, ticks=MaxNLocator(integer=True), **cbkwargs)
        ax.set_xlim(axparams[xtag]['lim'])
        ax.set_ylim(axparams[ytag]['lim'])
        pfit = Polyfit(x,y)
        fx = [-100,100]
        fy = []
        for ifx in fx:
            fy.append(pfit['poly'][0]*ifx+pfit['poly'][1])
        ax.plot(fx,fy)
        print pfit['det']
        ax.text(.05,.95, 'm = %3.3f, b = %3.3f, R^2 = %3.3f' % (pfit['poly'][0],pfit['poly'][1],pfit['det']),horizontalalignment='left',verticalalignment='top',transform=ax.transAxes)
        
    for i, (ax,subtag) in enumerate(zip(axes, TagParse(outfile))):
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            tick.label.set_rotation(45)
        for tick in ax.yaxis.get_major_ticks():
            tick.label.set_fontsize('small')
        #extent = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
        extent = matplotlib.transforms.Bbox.from_bounds(i*figx,0,figx,figy)
        plt.savefig(os.path.join(outdir, subtag+'.png'), bbox_inches=extent, transparent=True)
    plt.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight', transparent=True)
    plt.clf()

def TagParse(outfile):
    rawtags = outfile.split('.')[-1].split('-')
    subtags = []
    for dsol in rawtags[1].split('_'):
        for res in rawtags[2].split('_'):
            subtags.append('%s-%s-%s-%s'%(rawtags[0], dsol, res, rawtags[3]))
    return subtags

if __name__=="__main__":
    workdir=sys.argv[1]            
#    Unpickle(indir=workdir, infile='dsol_ele_hexbin', outdir=workdir, outfile='dsol_ele_hexbin', argn=2, plotter=Hexbin)
#    Unpickle(indir=workdir, infile='ratio_ddpk_hexbin', outdir=workdir, outfile='ratio_ddpk_hexbin', argn=2, plotter=Hexbin)
#    Unpickle(indir=workdir, infile='size_deepcount_hexbin', outdir=workdir, outfile='size_deepcount_hexbin', argn=2, plotter=Hexbin)
#    Unpickle(indir=workdir, infile='quick_dpk_full_dpk_multihexbin', outdir=workdir, outfile='quick_dpk_full_dpk_multihexbin', argn=6, plotter=MultiHexbin)
#    Unpickle(indir=workdir, infile='e4_dsol_e20_dsol_multihexbin', outdir=workdir, outfile='e4_dsol_e20_dsol_multihexbin', argn=6, plotter=MultiHexbin)
#    Unpickle(indir=workdir, infile='e4_ele_e20_ele_multihexbin', outdir=workdir, outfile='e4_ele_e20_ele_multihexbin', argn=6, plotter=MultiHexbin)
#    Unpickle(indir=workdir, infile='e4_ratio_e20_ratio_multihexbin', outdir=workdir, outfile='e4_ratio_e20_ratio_multihexbin', argn=6, plotter=MultiHexbin)
    #Unpickle(indir=workdir, infile='dsol_ele-alldsol-derk_h_y_snase-multihexbin', outdir=workdir, outfile='dsol_ele-alldsol-derk_h_y_snase-multihexbin', argn=8, plotter=MultiHexbinS)
    #Unpickle(indir=workdir, infile='ratio_ddpk-highdsol-derk_h_y_snase-multihexbin', outdir=workdir, outfile='ratio_ddpk-highdsol-derk_h_y_snase-multihexbin', argn=8, plotter=MultiHexbinS)
    #Unpickle(indir=workdir, infile='dpk_ddpk-alldsol-derk_his_tyr_snase-multihexbin', outdir=workdir, outfile='dpk_ddpk-alldsol-derk_his_tyr_snase-multihexbin', argn=8, plotter=MultiHexbinS)
    Unpickle(indir=workdir, infile='dpke4_dpke20-alldsol-derk_his_tyr_snase-multihexbin', outdir=workdir, outfile='dpke4_dpke20-alldsol-derk_his_tyr_snase-multihexbin', argn=8, plotter=MultiHexbinS)

